#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
客户端软件
实现与ESP32服务端的混合加密通信
"""

import socket
import json
import base64
import threading
import time
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.fernet import Fernet

class SecureClient:
    def __init__(self, server_host='localhost', server_port=8888):
        self.server_host = server_host
        self.server_port = server_port
        self.socket = None
        self.server_public_key = None
        self.fernet = None
        self.connected = False
        self.running = False
    
    def connect_to_server(self):
        """
        连接到服务器并建立加密通道
        """
        try:
            print(f"正在连接到服务器 {self.server_host}:{self.server_port}...")
            
            # 创建socket连接
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.server_host, self.server_port))
            
            print("TCP连接建立成功")
            
            # 1. 接收服务器公钥
            public_key_msg = self.receive_message()
            if public_key_msg and public_key_msg['type'] == 'public_key':
                public_key_pem = base64.b64decode(public_key_msg['data'])
                self.server_public_key = serialization.load_pem_public_key(public_key_pem)
                print("已接收服务器公钥")
            else:
                raise Exception("未能接收到服务器公钥")
            
            # 2. 生成AES密钥并用RSA公钥加密发送
            aes_key = Fernet.generate_key()
            self.fernet = Fernet(aes_key)
            
            # 使用服务器公钥加密AES密钥
            encrypted_aes_key = self.server_public_key.encrypt(
                aes_key,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            
            # 发送加密的AES密钥
            key_msg = {
                'type': 'encrypted_aes_key',
                'data': base64.b64encode(encrypted_aes_key).decode('utf-8')
            }
            self.send_message(key_msg)
            print("已发送加密的AES密钥")
            
            # 3. 接收握手确认
            confirm_msg = self.receive_message()
            if confirm_msg and confirm_msg['type'] == 'encrypted_data':
                # 解密确认消息
                encrypted_data = base64.b64decode(confirm_msg['data'])
                decrypted_data = self.fernet.decrypt(encrypted_data)
                handshake_response = json.loads(decrypted_data.decode('utf-8'))
                
                if handshake_response['type'] == 'handshake_complete':
                    print("加密通道建立成功！")
                    self.connected = True
                    return True
            
            raise Exception("握手失败")
            
        except Exception as e:
            print(f"连接失败：{e}")
            if self.socket:
                self.socket.close()
            return False
    
    def send_message(self, message):
        """
        发送JSON消息
        """
        try:
            data = json.dumps(message).encode('utf-8')
            length = len(data)
            self.socket.sendall(length.to_bytes(4, byteorder='big'))
            self.socket.sendall(data)
        except Exception as e:
            print(f"发送消息失败：{e}")
    
    def receive_message(self):
        """
        接收JSON消息
        """
        try:
            # 先接收消息长度
            length_bytes = self.socket.recv(4)
            if not length_bytes:
                return None
            
            length = int.from_bytes(length_bytes, byteorder='big')
            
            # 接收消息数据
            data = b''
            while len(data) < length:
                chunk = self.socket.recv(length - len(data))
                if not chunk:
                    return None
                data += chunk
            
            return json.loads(data.decode('utf-8'))
        except Exception as e:
            print(f"接收消息失败：{e}")
            return None
    
    def send_encrypted_message(self, message):
        """
        发送加密消息
        """
        if not self.connected or not self.fernet:
            print("错误：未建立加密连接")
            return False
        
        try:
            # 加密消息
            message_json = json.dumps(message).encode('utf-8')
            encrypted_data = self.fernet.encrypt(message_json)
            
            # 发送加密消息
            encrypted_msg = {
                'type': 'encrypted_data',
                'data': base64.b64encode(encrypted_data).decode('utf-8')
            }
            self.send_message(encrypted_msg)
            return True
        except Exception as e:
            print(f"发送加密消息失败：{e}")
            return False
    
    def receive_encrypted_message(self):
        """
        接收并解密消息
        """
        if not self.connected or not self.fernet:
            print("错误：未建立加密连接")
            return None
        
        try:
            # 接收加密消息
            encrypted_msg = self.receive_message()
            if not encrypted_msg or encrypted_msg['type'] != 'encrypted_data':
                return None
            
            # 解密消息
            encrypted_data = base64.b64decode(encrypted_msg['data'])
            decrypted_data = self.fernet.decrypt(encrypted_data)
            message = json.loads(decrypted_data.decode('utf-8'))
            
            return message
        except Exception as e:
            print(f"接收加密消息失败：{e}")
            return None
    
    def ping_server(self):
        """
        发送ping消息测试连接
        """
        ping_msg = {
            'type': 'ping',
            'timestamp': time.time()
        }
        
        if self.send_encrypted_message(ping_msg):
            print("已发送ping消息")
            response = self.receive_encrypted_message()
            if response:
                print(f"收到服务器响应：{response}")
                return True
        return False
    
    def request_sensor_data(self):
        """
        请求传感器数据
        """
        request_msg = {
            'type': 'sensor_request',
            'timestamp': time.time()
        }
        
        if self.send_encrypted_message(request_msg):
            print("已请求传感器数据")
            response = self.receive_encrypted_message()
            if response:
                print(f"传感器数据：{response}")
                return response
        return None
    
    def send_control_command(self, action):
        """
        发送控制命令
        """
        control_msg = {
            'type': 'control',
            'action': action,
            'timestamp': time.time()
        }
        
        if self.send_encrypted_message(control_msg):
            print(f"已发送控制命令：{action}")
            response = self.receive_encrypted_message()
            if response:
                print(f"控制响应：{response}")
                return response
        return None
    
    def start_interactive_mode(self):
        """
        启动交互模式
        """
        print("\n=== 交互模式 ===")
        print("可用命令：")
        print("1. ping - 测试连接")
        print("2. sensor - 获取传感器数据")
        print("3. control <action> - 发送控制命令")
        print("4. quit - 退出")
        print()
        
        while self.connected:
            try:
                command = input("请输入命令: ").strip().split()
                if not command:
                    continue
                
                cmd = command[0].lower()
                
                if cmd == 'quit':
                    break
                elif cmd == 'ping':
                    self.ping_server()
                elif cmd == 'sensor':
                    self.request_sensor_data()
                elif cmd == 'control':
                    if len(command) > 1:
                        action = ' '.join(command[1:])
                        self.send_control_command(action)
                    else:
                        print("请指定控制动作，例如：control turn_on_led")
                else:
                    print(f"未知命令：{cmd}")
                    
            except KeyboardInterrupt:
                break
            except Exception as e:
                print(f"命令执行出错：{e}")
    
    def disconnect(self):
        """
        断开连接
        """
        self.connected = False
        if self.socket:
            self.socket.close()
        print("已断开连接")

def main():
    print("=== 安全客户端 ===")
    
    client = SecureClient()
    
    try:
        # 连接到服务器
        if client.connect_to_server():
            # 启动交互模式
            client.start_interactive_mode()
        else:
            print("连接失败，程序退出")
    except KeyboardInterrupt:
        print("\n收到中断信号，正在断开连接...")
    finally:
        client.disconnect()

if __name__ == "__main__":
    main()